from pgmpy.factors.discrete import TabularCPD
from pgmpy.inference import VariableElimination
from pgmpy.models import DiscreteBayesianNetwork

# 创建离散贝叶斯网络模型
model = DiscreteBayesianNetwork([
    ("Rain", "Maintenance"),
    ("Maintenance", "Train"),
    ("Rain", "Train"),
    ("Train", "Appointment")
])

# 定义条件概率分布字典
# 当变量有m种取值，父节点有n种取值时，矩阵形状为 m * n
# 每列的概率之和应为1
cpd_dict = {
    "Rain": TabularCPD(
        "Rain", 3,
        [[0.7], [0.2], [0.1]],
        state_names={"Rain": ["无雨", "小雨", "大雨"]}
    ),

    "Maintenance": TabularCPD(
        "Maintenance", 2,
        [
            [0.4, 0.2, 0.1],
            [0.6, 0.8, 0.9]
        ],
        evidence=['Rain'],
        evidence_card=[3],
        state_names={
            "Rain": ["无雨", "小雨", "大雨"],
            "Maintenance": ["是", "否"]
        }
    ),

    "Train": TabularCPD(
        "Train", 2,
        [
            [0.8, 0.9, 0.6, 0.7, 0.4, 0.5],
            [0.2, 0.1, 0.4, 0.3, 0.6, 0.5]
        ],
        evidence=["Rain", "Maintenance"],
        evidence_card=[3, 2],
        state_names={
            "Rain": ["无雨", "小雨", "大雨"],
            "Maintenance": ["是", "否"],
            "Train": ["准点", "延误"]
        }
    ),

    "Appointment": TabularCPD(
        "Appointment", 2,
        [[0.9, 0.6], [0.1, 0.4]],
        evidence=["Train"],
        evidence_card=[2],
        state_names={
            "Train": ["准点", "延误"],
            "Appointment": ["出席", "缺席"]
        }
    )
}

# 将CPD添加到模型中
for val in cpd_dict.values():
    model.add_cpds(val)

# 检查模型是否有效
print("通过检测" if model.check_model() else "未通过")

# 进行推理计算
infer = VariableElimination(model)

# 查询所有变量的联合概率分布
p1 = infer.query(variables=["Rain", "Maintenance", "Train", "Appointment"])
print(p1)

# 在给定证据条件下查询变量概率
p2 = infer.query(variables=["Appointment"], evidence={"Train": "延误"})
print(p2)
